from sklearn.linear_model import LinearRegression
from sklearn.metrics import accuracy_score, r2_score
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

from data import DataUtils
from visual import VisualUtils


class TrainUtils:
	lr = None
	tr = None

	# 训练线性回归模型
	@staticmethod
	def train_linear_model(X, y):
		TrainUtils.lr = LinearRegression().fit(X, y)

	# 训练决策树模型
	@staticmethod
	def train_tree_model(X, y):
		# max_depth设置树深度
		TrainUtils.tr = DecisionTreeClassifier(min_samples_split=2, min_samples_leaf=10, criterion='entropy', splitter='random').fit(X, y)
		view = VisualUtils()
		view.plt_tree(TrainUtils.tr)

	# 模型预测
	@staticmethod
	def linear_model_predict(test_x, test_y):
		if TrainUtils.lr is not None:
			view = VisualUtils()
			predict = TrainUtils.lr.predict(test_x)
			r2 = r2_score(test_y, predict, sample_weight=None)
			view.plt_linear(test_y, predict, r2)
			return predict, r2
		else:
			return None, None

	# 模型预测
	@staticmethod
	def tree_model_predict(test_x, test_y):
		if TrainUtils.tr is not None:
			predict = TrainUtils.tr.predict(test_x)
			score = accuracy_score(test_y, predict, sample_weight=None)
			view = VisualUtils()
			view.plot_confusion_matrix(test_y, predict, "决策树-混淆矩阵 score = {}".format(score))
			view.plt_pie(TrainUtils.tr.feature_importances_, DataUtils.names[:8], "决策树-特征值权重占比")
			return predict, score
		else:
			return None, None

	@staticmethod
	def default_predict(X, y):
		train_x, test_x, train_y, test_y = train_test_split(X, y, test_size=0.1, random_state=0)
		TrainUtils.tr.fit(train_x, train_y)
		tr_pred_y = TrainUtils.tr.predict(test_x)
		view = VisualUtils()
		view.plot_confusion_matrix(test_y, tr_pred_y, "决策树-混淆矩阵 score = {}".format(accuracy_score(test_y, tr_pred_y, sample_weight=None)))
		view.plt_pie(TrainUtils.tr.feature_importances_, DataUtils.names[:8], "特征值权重占比")

		lr_pred_y = TrainUtils.lr.predict(test_x)
		view.plt_linear(test_y, lr_pred_y, r2_score(test_y, lr_pred_y, sample_weight=None))
